import os.path as osp
from glob import glob

import torchvision
from PIL import Image


class NewCIFAR10(torchvision.datasets.CIFAR10):
    def __init__(
            self, root, extra_imgs_root,
            train=True, transform=None,
            target_transform=None, download=False
    ):
        super().__init__(root, train, transform, target_transform, download)
        self.extra_imgs_root = extra_imgs_root
        self.extra_imgs = sorted(glob(self.extra_imgs_root + "/*.jpg"))
        total = len(self.extra_imgs)
        split_pos = int(0.7 * total)
        if self.train:
            self.extra_imgs = self.extra_imgs[:split_pos]
        else:
            self.extra_imgs = self.extra_imgs[split_pos:]
        self.extra_imgs = list(map(
            lambda x: Image.open(x), self.extra_imgs
        ))

    def __getitem__(self, index):
        super_len = super(NewCIFAR10, self).__len__()
        if index < super_len:
            return super(NewCIFAR10, self).__getitem__(index)
        else:
            img = self.extra_imgs[index - super_len]
            target = 10

            if self.transform is not None:
                img = self.transform(img)

            if self.target_transform is not None:
                target = self.target_transform(target)

            return img, target

    def __len__(self):
        return super(NewCIFAR10, self).__len__() + len(self.extra_imgs)
